#!/usr/bin/env python
# -*- coding: utf-8 -*-

import json
import yaml
import logging
import subprocess

# Find the best implementation available on this platform
try:
    from cStringIO import StringIO #@UnusedImport
except:
    from StringIO import StringIO  #@Reimport

__all__ = ['submit', 'status']

log = logging.getLogger("glite.wms.job")

def submit(jdl_path, endpoint=None):
    """
    Execute the glite-wms-job-submit command and return the job ID.
    Example: https://wms04.pic.es:9000/kL_UMwt4ktZjKirzErrq9w
    
    @type  endpoint: URL string
    @param endpoint: URL of the WMS service to be used.
    @type  jdl_path: string
    @param jdl_path: Path to the JDL file to be submitted.
    
    @rtype:  string
    @return: An URL string representing the job ID.
    """
    
    command = ['glite-wms-job-submit', '-a', '--json']
    
    if endpoint:
        command.extend(('-e', endpoint))
    
    command.append(jdl_path)
    
    log.debug("Submitting WMS job:\n%s", " ".join(command))
    p = subprocess.Popen(command, shell=False, stdin  = None,
                                               stdout = subprocess.PIPE,
                                               stderr = subprocess.PIPE)
    (out, err) = p.communicate()
    log.debug("Command finished.")
    
    p.stdout.close()
    p.stderr.close()
    
    if p.returncode > 0:
        msg = "glite-wms-job-submit failed with return code %s. Error message:\n%s" % (p.returncode, err)
        log.error(msg)
        raise RuntimeError(msg)
    
    if p.returncode < 0:
        msg = "glite-wms-job-submit was killed by signal %d." % (-p.returncode)
        log.error(msg)
        raise RuntimeError(msg)
    
    status = json.loads(out)
    
    if status['result'] != 'success':
        msg = "glite-wms-job-submit was unable to complete the operation. Error message:\n%s" % (err)
        log.error(msg)
        raise RuntimeError(msg)
    
    log.debug("Submitted WMS job with id %s.", status['jobid'])
    return status['jobid']

def status(jobids):
    """
    Execute the glite-wms-job-status command and return the job status information.
    
    @type  jobids: string or tuple of strings.
    @param jobid: The list of jobids to query its status.
    
    @rtype:  dict
    @return: A mapping structure containing the status of each job requested.
    """
    
    command = ['glite-wms-job-status', '--verbosity', '2']
    
    if isinstance(jobids, basestring):
        jobids = (jobids,)
    
    command.extend(jobids)
    
    log.debug("Checking WMS job status:\n%s", " ".join(command))
    p = subprocess.Popen(command, shell=False, stdin  = None,
                                               stdout = subprocess.PIPE,
                                               stderr = subprocess.PIPE)
    
    (out, err) = p.communicate()
    log.debug("Command finished.")
    
    p.stdout.close()
    p.stderr.close()
    
    if p.returncode > 0:
        msg = "glite-wms-job-status failed with return code %s («%s»)." % (p.returncode, err)
        log.error(msg)
        raise RuntimeError(msg)
    
    if p.returncode < 0:
        msg = "glite-wms-job-status was killed by signal %d." % (-p.returncode)
        log.error(msg)
        raise RuntimeError(msg)
    
    status = {}
    job_status = None
    current = StringIO()
    for line in out.splitlines():
        line = line.strip()
        if not line or \
          line == "======================= glite-wms-job-status Success =====================" or \
          line == "BOOKKEEPING INFORMATION:" :
            continue
        if line == '---':
            job_status = yaml.safe_load(current.getvalue().replace('    - ', ' - |\n '))
            status[job_status['Status info for the Job']] = job_status
            del job_status['Status info for the Job']
            current = StringIO()
            current.write('\n')
            continue
        if line == "==========================================================================" :
            job_status.update(yaml.safe_load(current.getvalue().replace(' =    ', ' : |\n ').replace('\n- ', '\n')))
            current = StringIO()
            continue
        current.write(line)
        current.write('\n')
    
    return status

def output(jobids, output_dir=None):
    """
    Execute the glite-wms-job-output command and return the command status.
    
    @type  jobids: string or tuple of strings.
    @param jobid: The list of jobids to get its output.
    @type  output_dir: string
    @param output_dir: Path to store the outputs into.
    
    @rtype:  dict
    @return: A mapping structure containing the status of each output requested.
    """
    
    command = ['glite-wms-job-output', '--json']
    
    if output_dir:
        command.extend(('--dir', output_dir))
    
    if isinstance(jobids, basestring):
        jobids = (jobids,)
    
    command.extend(jobids)
    
    log.debug("Retrieving WMS job output:\n%s", " ".join(command))
    p = subprocess.Popen(command, shell=False, stdin  = None,
                                               stdout = subprocess.PIPE,
                                               stderr = subprocess.PIPE)
    (out, err) = p.communicate()
    log.debug("Command finished.")
    
    p.stdout.close()
    p.stderr.close()
    
    if p.returncode > 0:
        msg = "glite-wms-job-output failed with return code %s. Error message:\n%s" % (p.returncode, err)
        log.error(msg)
        raise RuntimeError(msg)
    
    if p.returncode < 0:
        msg = "glite-wms-job-output was killed by signal %d." % (-p.returncode)
        log.error(msg)
        raise RuntimeError(msg)
    
    status = json.loads(out)
    
    if status['result'] != 'success':
        msg = "glite-wms-job-output was unable to complete the operation. Error message:\n%s" % (err)
        log.error(msg)
        raise RuntimeError(msg)
    
    jobs = status['jobs']
    
    if len(jobs) == 1:
        return jobs.popitem()[1]
    
    return jobs
